import numpy as np
import scipy.spatial.distance
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from utils2 import JS_D
import time

ds_path = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/round2_2ds.npy'
ds_feat = np.load(ds_path)

frame_count = np.shape(ds_feat)[0]

SR_typicalScene = ds_feat[408]
TR_typicalScene = ds_feat[756]
AT_typicalScene = ds_feat[25182]

SR_scp = []
TR_scp = []
AT_scp = []

for i in range(frame_count):
    cur_ds = ds_feat[i]
    js_s = JS_D(cur_ds, SR_typicalScene)
    js_t = JS_D(cur_ds, TR_typicalScene)
    js_a = JS_D(cur_ds, AT_typicalScene)
    if js_s < js_t and js_s < js_a:
        SR_scp.append(cur_ds)
    elif js_t < js_s and js_t < js_a:
        TR_scp.append(cur_ds)
    else:
        AT_scp.append(cur_ds)

SR_scp = np.average(np.array(SR_scp), axis=0)
TR_scp =  np.average(np.array(TR_scp), axis=0)
AT_scp =  np.average(np.array(AT_scp), axis=0)



# scp_path = 'E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/ds_cluster_eucliden/centre8.npy'
# centre = np.load(scp_path)
# print(JS_D(SR_scp, centre[0]))
# print(JS_D(TR_scp, centre[1]))
# print(JS_D(AT_scp, centre[2]))

fig = plt.figure(figsize=(25,12))
ax = fig.add_subplot(1, 3, 1)
plt.bar(range(100), SR_scp,width=1, align='edge')
plt.ylim(0,0.08)
ax.set_xticks([0,25,50,75,99])
ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

ax = fig.add_subplot(1, 3, 2)
plt.bar(range(100), TR_scp,width=1, align='edge')
plt.ylim(0,0.08)
ax.set_xticks([0,25,50,75,99])
ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

ax = fig.add_subplot(1, 3, 3)
plt.bar(range(100), AT_scp,width=1, align='edge')
plt.ylim(0,0.08)
ax.set_xticks([0,25,50,75,99])
ax.set_xticklabels(['-1','-0.5','0','0.5','1'])

scp = np.array((SR_scp, TR_scp, AT_scp))
np.save('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/scp_byTypicalScene.npy', scp)
plt.savefig('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/scp_byTypicalScene.png')

# plt.savefig('E:/Research/2020ContrastiveLearningForSceneLabel/Data/20210329ExperimentData/round2/ds_cluster_eucliden/centre%d.png'%i)
plt.show()
plt.close()